package websocket

import (
	"encoding/json"
	"sync"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
)

// Wrapper WebSocket 通用封装类
type Wrapper struct {
	hub       *Hub
	upgrader  websocket.Upgrader
	hubsByID  map[string]*Hub // 支持多个 Hub（按命名空间或业务区分）
	hubsMutex sync.RWMutex
	rooms     map[string]*Room // 房间管理（房间ID -> Room）
	roomsMutex sync.RWMutex
}

// NewWrapper 创建 WebSocket 封装实例
func NewWrapper() *Wrapper {
	w := &Wrapper{
		hub:       NewHub(),
		upgrader:  DefaultUpgrader,
		hubsByID:  make(map[string]*Hub),
		rooms:     make(map[string]*Room),
	}

	// 启动默认 Hub
	go w.hub.Run()

	return w
}

// GetHub 获取默认 Hub
func (w *Wrapper) GetHub() *Hub {
	return w.hub
}

// GetOrCreateHub 获取或创建指定 ID 的 Hub
func (w *Wrapper) GetOrCreateHub(id string) *Hub {
	w.hubsMutex.Lock()
	defer w.hubsMutex.Unlock()

	if hub, ok := w.hubsByID[id]; ok {
		return hub
	}

	hub := NewHub()
	w.hubsByID[id] = hub
	go hub.Run()

	return hub
}

// Upgrade 升级 HTTP 连接为 WebSocket
func (w *Wrapper) Upgrade(c *gin.Context) (*Connection, error) {
	conn, err := w.upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		return nil, err
	}

	return NewConnection(conn), nil
}

// HandleConnection 处理 WebSocket 连接（使用默认 Hub）
func (w *Wrapper) HandleConnection(
	c *gin.Context,
	connectionID string,
	onConnect func(*Connection),
	onMessage func(*Connection, []byte) error,
	onDisconnect func(*Connection),
) error {
	conn, err := w.Upgrade(c)
	if err != nil {
		return err
	}

	// 注册连接
	w.hub.RegisterConnection(connectionID, conn)

	// 连接回调
	if onConnect != nil {
		onConnect(conn)
	}

	// 启动读写 goroutine
	go conn.WritePump()
	
	// 启动读取 goroutine，读取结束后触发断开回调
	go func() {
		conn.ReadPump(func(message []byte) error {
			if onMessage != nil {
				return onMessage(conn, message)
			}
			return nil
		})
		// ReadPump 结束意味着连接关闭
		if onDisconnect != nil {
			onDisconnect(conn)
		}
		w.hub.UnregisterConnection(connectionID, conn)
	}()

	return nil
}

// Broadcast 广播消息到所有连接（默认 Hub）
func (w *Wrapper) Broadcast(message []byte) {
	w.hub.Broadcast(message)
}

// BroadcastJSON 广播 JSON 消息
func (w *Wrapper) BroadcastJSON(v interface{}) error {
	data, err := json.Marshal(v)
	if err != nil {
		return err
	}
	w.hub.Broadcast(data)
	return nil
}

// BroadcastMessage 广播结构化消息
func (w *Wrapper) BroadcastMessage(msgType string, payload interface{}) error {
	msg := Message{
		Type:    msgType,
		Payload: payload,
	}
	return w.BroadcastJSON(msg)
}

// SendToID 发送消息到指定连接ID（默认 Hub）
func (w *Wrapper) SendToID(id string, message []byte) error {
	return w.hub.SendToID(id, message)
}

// SendJSONToID 发送 JSON 消息到指定连接ID
func (w *Wrapper) SendJSONToID(id string, v interface{}) error {
	data, err := json.Marshal(v)
	if err != nil {
		return err
	}
	return w.hub.SendToID(id, data)
}

// SendMessageToID 发送结构化消息到指定连接ID
func (w *Wrapper) SendMessageToID(id string, msgType string, payload interface{}) error {
	msg := Message{
		Type:    msgType,
		Payload: payload,
	}
	return w.SendJSONToID(id, msg)
}

// GetConnectionCount 获取连接数（默认 Hub）
func (w *Wrapper) GetConnectionCount() int {
	return w.hub.GetConnectionCount()
}

// GetConnections 获取所有连接ID（默认 Hub）
func (w *Wrapper) GetConnections() []string {
	return w.hub.GetConnections()
}

// SetUpgrader 设置自定义 Upgrader
func (w *Wrapper) SetUpgrader(upgrader websocket.Upgrader) {
	w.upgrader = upgrader
}

// Close 关闭所有连接
func (w *Wrapper) Close() error {
	// 关闭所有 Hub
	w.hubsMutex.Lock()
	defer w.hubsMutex.Unlock()

	for _, hub := range w.hubsByID {
		hub.mu.Lock()
		for conn := range hub.connections {
			conn.Close()
		}
		hub.mu.Unlock()
	}

	// 关闭默认 Hub
	w.hub.mu.Lock()
	for conn := range w.hub.connections {
		conn.Close()
	}
	w.hub.mu.Unlock()

	return nil
}

